import os
import time
import tensorflow as tf
import pandas as pd
import matplotlib.pyplot as plt
from tensorflow.keras.models import save_model

from imagenette.Utils.data_generator import train_val_generator
from imagenette.Utils.image_plot import plot_image
from imagenette.Nets.Conv_net import Conv_model

# 读取训练集
train_gen = train_val_generator(
    data_dir='../dataset/train',
    target_size=(160, 160),
    batch_size=32,
    class_mode='categorical',
    subset='training'
)

# 读取验证集
val_gen = train_val_generator(
    data_dir='../dataset/train',
    target_size=(160, 160),
    batch_size=32,
    class_mode='categorical',
    subset='validation'
)

# 从训练集读取标签列表
class_names = list(train_gen.class_indices.keys())
print(class_names)

# 显示训练集的15张图片
train_data, train_labels = train_gen.next()
plot_image(train_data, train_labels, class_names)
print(len(train_gen))

# 显示验证集的15张图片
val_data, val_labels = val_gen.next()
plot_image(val_data, val_labels, class_names)

# 编译模型
model = Conv_model()
model.compile(
    # loss='categorical_crossentropy',
    # optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
    loss='kullback_leibler_divergence',
    # loss='mean_absolute_percentage_error',
    optimizer=tf.keras.optimizers.Adam(),
    metrics=['accuracy']
)

# 训练模型
history = model.fit(
    x=train_gen,
    steps_per_epoch=174,
    epochs=50,
    validation_data=val_gen,
    validation_steps=44,
    shuffle=True
)

# 输出训练结果
pd.DataFrame(history.history).plot(figsize=(8, 5))
plt.grid(True)
plt.xlabel('epoch')
plt.show()

# 保存模型
model_name = "my_model_" + time.strftime('%Y-%m-%d')
model_path = os.path.join('..', 'Models', model_name)
if not os.path.exists(model_path):
    os.makedirs(model_path)

save_model(model=model, filepath=model_path)